import argparse
import signal
import traceback

from engine_server.ctlapp.ctl import Ctler
from engine_server.engine.engine import BaseEngineMgr
from engine_server.utils.logger import run_log


def setup_signal_handlers(engine: BaseEngineMgr, ctl: Ctler):
    def handle_shutdown(signum):
        print(f"Received signal {signum}, starting shutdown...")
        ctl.stop()
        engine.stop()

    signal.signal(signal.SIGINT, handle_shutdown)
    signal.signal(signal.SIGTERM, handle_shutdown)


def main():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('--es-engine', type=str, required=True, choices=['vllm', 'sglang', 'mindie'])
    known_args, remaining_args = parser.parse_known_args()

    if known_args.es_engine not in ['vllm']:
        raise RuntimeError(f"Unsupported engine: {known_args.es_engine}")
    else:
        ctl = Ctler()
        ctl.start()
        from engine_server.engine.vllm.engine import VLLMEngineMgr
        engine = VLLMEngineMgr(remaining_args)
        setup_signal_handlers(engine, ctl)
        try:
            engine.start()
        except Exception as e:
            run_log.error(f"Engine get exception: {e}")
            traceback.print_exc()
        finally:
            ctl.stop()
            engine.stop()
            run_log.info("Engine and Ctler stopped")


if __name__ == '__main__':
    main()